import requests
import json
from openai import OpenAI
from llm.qianwen_prompt_steps import *

def qianwen_with_step(code: str="你好", data_dict: str="", initial_code:str="你好", step: int=0):
    try:
        client = OpenAI(
            # api_key="sk-3f966520a2684aea9367de0ebe63508e",
            api_key="sk-0a647c7ea8af4d29a704add3c6efb125",
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

        # qianwen_prompt = qianwen_prompt_steps[step].replace("{{ input_code }}", code).replace("{{ input_dict }}", data_dict).replace("{{ raw_input_code }}", initial_code)
        # print(f"==================== step{step + 1} prompt ====================")
        # print(qianwen_prompt)

        completion = client.chat.completions.create(
            # model="qwen2.5-72b-instruct",
            model="deepseek-v3",
            messages=[
                {
                    'role': 'system', 
                    'content': qianwen_system_prompt
                },
                {
                    'role': 'user', 
                    'content': qianwen_prompt_steps[step].replace("{{ input_code }}", code).replace("{{ input_dict }}", data_dict).replace("{{ raw_input_code }}", initial_code)
                }
            ],
            temperature=0,
            top_p=0.1,
        )
        response_data = json.loads(completion.model_dump_json())
        content = response_data["choices"][0]["message"]["content"]
        return content
    except Exception as e:
        print(e)

def model_with_step(url: str="", model: str="", code: str="你好", data_dict: str="", initial_code:str="你好", step: int=0):
    payload = {
      "model": model,
      "prompt": qianwen_prompt_steps[step].replace("{{ input_code }}",code).replace("{{ input_dict }}", data_dict).replace("{{ raw_input_code }}", initial_code),
      "stream": False,
    }
    try:
          response = requests.post(url=url, json=payload)
          response.raise_for_status()
          content = json.loads(response.text)
          return content["response"]
    except Exception as e:
          print(e)

code = """IF1(j - m >= 0)
    IF2(m > 10)
        tdyn = 0;
    ENDIF2
    m++;
ELSE1
    tdyn = 0;
ENDIF1
"""

if __name__ == '__main__':
    print(qianwen_with_step(code, "", "", 0))